import os
import csv

import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
from scipy.constants import physical_constants

WORKSPACE = os.getcwd()
DETECTOR_H = 4
DETECTOR_R = 2
DETECTOR_D = 20
HOLE_D = 1
SOURCE_ENERGY = 0.662
E0 = physical_constants["natural unit of energy in MeV"][0]

class FinalDetector:
    sections = {
        "energy": list(),
        "photoelectric": list(),
        "compton": list()
    }
    
    def __init__(self) -> None:
        with open(os.path.join(WORKSPACE, "data.csv"), "r") as dataFile:
            csvReader = csv.reader(dataFile)
            for row in csvReader:
                self.sections["energy"].append(float(row[0]))
                self.sections["photoelectric"].append(float(row[1]))
                self.sections["compton"].append(float(row[2]))
        self.photoelectricSectionCalculator = interp1d(self.sections["energy"], self.sections["photoelectric"])
        self.comptonSectionCalculator = interp1d(self.sections["energy"], self.sections["compton"])
                
    def simulate(self, particleNumber=1e6):
        particleNumber = int(particleNumber)
        
        kesai0 = np.random.rand(particleNumber)
        kesai1 = np.random.rand(particleNumber)
        particles_vx = np.sin(np.pi * kesai0) * np.cos(2 * np.pi * kesai1)
        particles_vy = np.sin(np.pi * kesai0) * np.sin(2 * np.pi * kesai1)
        particles_vz = np.cos(np.pi * kesai0)
        particles_vr = np.sqrt(np.square(particles_vx) + np.square(particles_vy))
        particles_v = np.sqrt(np.square(particles_vr) + np.square(particles_vz))
        
        particles_energy = np.ones(particleNumber) * SOURCE_ENERGY
        particles_alive = np.ones(particleNumber)
        particles_detected = np.zeros(particleNumber)
        
        particles_x = DETECTOR_D * particles_vx / np.abs(particles_vz)
        particles_y = DETECTOR_D * particles_vy / np.abs(particles_vz)
        particles_r = np.sqrt(np.square(particles_x) + np.square(particles_y))
        particles_alive[particles_r > HOLE_D / 2] = 0
        particles_z = DETECTOR_H * np.ones(particleNumber) + DETECTOR_D * (1 + np.sign(particles_vz))
        particles_alive[particles_z > DETECTOR_H] = 0        
        particles_alive_init = particles_alive.copy()
        
        particles_photoelectric_section = np.zeros(particleNumber)
        particles_compton_section = np.zeros(particleNumber)
        particles_section = particles_photoelectric_section + particles_compton_section
        particles_dl = np.zeros(particleNumber)
        
        while np.sum(particles_alive) > 0:
            particles_alive[particles_energy < np.min(self.sections["energy"])] = 0
            particles_photoelectric_section[particles_alive == 1], particles_compton_section[particles_alive == 1] = self.calculateSection(particles_energy[particles_alive == 1])
            particles_section = particles_photoelectric_section + particles_compton_section
            
            particles_dl[particles_alive == 1] = np.random.exponential(1 / (particles_section[particles_alive == 1]))
            particles_vr = np.sqrt(np.square(particles_vx) + np.square(particles_vy))
            particles_v = np.sqrt(np.square(particles_vr) + np.square(particles_vz))
            particles_x[particles_alive == 1] += particles_dl[particles_alive == 1] * particles_vx[particles_alive == 1] / particles_v[particles_alive == 1]
            particles_y[particles_alive == 1] += particles_dl[particles_alive == 1] * particles_vy[particles_alive == 1] / particles_v[particles_alive == 1]
            particles_z[particles_alive == 1] += particles_dl[particles_alive == 1] * particles_vz[particles_alive == 1] / particles_v[particles_alive == 1]
            particles_r = np.sqrt(np.square(particles_x) + np.square(particles_y))
            particles_alive[particles_r > DETECTOR_R] = 0
            particles_alive[particles_z < 0] = 0
            particles_alive[particles_z > DETECTOR_H] = 0
            
            kesai = np.random.rand(particleNumber)
            particles_alive_copy = particles_alive.copy()
            particles_alive[particles_alive_copy * kesai * particles_section > particles_compton_section] = 0
            particles_energy[particles_alive_copy * kesai * particles_section > particles_compton_section] = 0
            particles_detected[particles_alive_copy * kesai * particles_section > particles_compton_section] = 1
            
            particles_energy, particles_vx, particles_vy, particles_vz, particles_detected = self.simulateCompton(particles_alive, particles_energy, particles_vx, particles_vy, particles_vz, particles_detected)
        
        particles_energy_despoit = (np.ones(particleNumber) * SOURCE_ENERGY - particles_energy)[particles_detected == 1]
        self.plotData(particles_energy_despoit, np.sum(particles_alive_init), np.sum(particles_detected))
                
    def plotData(self, energy: np.ndarray, particlesN: int, particlesNm: int) -> None:
        efficiency = particlesNm / particlesN
        efficiency_sigma = efficiency * (1 - efficiency)
        efficiency_delta = 2 * efficiency_sigma * np.sqrt(particlesN) / particlesNm
        particles_peak = np.zeros(len(energy))
        histResult, bins = np.histogram(energy, 160, (0, 0.8), False)
        energyMax = np.average((bins[np.argmax(histResult)], bins[np.argmax(histResult) + 1]))
        energy = self.addNoise(energy)
        histResult, bins, patches = plt.hist(energy, 160, (0, 0.8), False, label="total\n${\eta}_{n}=%.3f$\n${\sigma}_{\eta}=%.3f$\n${\delta}_{\eta}=%.3e$"%(efficiency, efficiency_sigma, efficiency_delta))
        FWHM = 0.01 + 0.05 * np.sqrt(energyMax + 0.4 * np.square(energyMax))
        sigma = FWHM / (2 * np.sqrt(2 * np.log(2)))
        energyLower = energyMax - 3 * sigma
        particles_peak[energy >= energyLower] = 1
        particlesNp = np.sum(particles_peak)
        ratio = particlesNp / particlesNm
        ratio_sigma = ratio * (1 - ratio)
        ratio_delta = 2 * ratio_sigma * np.sqrt(particlesNm) / particlesNp
        binsCalculator = interp1d(bins, bins, kind="nearest")
        binsLower = binsCalculator(energyLower)
        plt.hist(energy[energy >= binsLower], bins[bins >= binsLower], (0, 0.8), False, label="peak\n${\eta}_{p}=%.3f$\n${\sigma}_{\eta p}=%.3f$\n${\delta}_{\eta p}=%.3e$"%(ratio, ratio_sigma, ratio_delta))
        binsNew = (bins[1:] + bins[:-1]) / 2
        binsFWHM = binsNew[histResult >= max(histResult) / 2]
        binsFWHMMax = max(binsFWHM)
        binsFWHMMin = min(binsFWHM)
        FWHM = binsFWHMMax - binsFWHMMin
        plt.plot([binsFWHMMin, binsFWHMMin], [0, max(histResult) / 2], label="half peak (low)\n$E=%.3fMeV$"%binsFWHMMin)
        plt.plot([energyMax, energyMax], [0, max(histResult)], label="peak\n$E=%.3fMeV$\n$\\alpha=%.2e$"%(energyMax, FWHM / energyMax))
        plt.plot([binsFWHMMax, binsFWHMMax], [0, max(histResult) / 2], label="half peak (high)\n$E=%.3fMeV$"%binsFWHMMax)
        plt.grid()
        plt.legend()
        plt.savefig(os.path.join(WORKSPACE, "result.png"))
        
    def addNoise(self, energy: np.ndarray) -> np.ndarray:
        FWHM = 0.01 + 0.05 * np.sqrt(energy + 0.4 * np.square(energy))
        sigma = FWHM / (2 * np.sqrt(2 * np.log(2)))
        return energy + sigma * np.random.standard_normal(len(energy))
        
    def calculateSection(self, energy: np.ndarray) -> tuple:
        return self.photoelectricSectionCalculator(energy), self.comptonSectionCalculator(energy)
    
    def simulateCompton(self, aliveState: np.ndarray, particlesEnergy: np.ndarray, particles_vx: np.ndarray, particles_vy: np.ndarray, particles_vz: np.ndarray, particles_detected: np.ndarray) -> tuple:
        for index, state in enumerate(aliveState):
            if state == 1:
                alpha = particlesEnergy[index] / E0
                x_ = 0
                while True:
                    kesai0 = np.random.rand()
                    kesai1 = np.random.rand()
                    kesai2 = np.random.rand()
                    if kesai0 <= 27 / (4 * alpha + 29):
                        x1 = (1 + 2 * alpha) / (1 + 2 * alpha * kesai1)
                        if kesai2 <= 0.5 * (np.square((alpha + 1 - x1) / alpha) + 1):
                            x_ = x1
                            break
                    else:
                        x2 = 1 + 2 * alpha * kesai1
                        if kesai2 <= 27 / 4 * np.square(x2 - 1) / np.power(x2, 3):
                            x_ = x2
                            break
                particles_detected[index] = 1
                particlesEnergy[index] /= x_
                alpha_ = particlesEnergy[index] / E0
                a = 1 - 1 / alpha_ + 1 / alpha
                b = np.sqrt(1 - np.square(a))
                kesai_ = 2 * np.pi * np.random.rand()
                c = np.cos(kesai_)
                d = np.sin(kesai_)
                v_ = np.sqrt(np.square(particles_vx[index]) + np.square(particles_vx[index]))
                if v_ == 0:
                    particles_vx[index] = b * c
                    particles_vy[index] = b * d
                    particles_vz[index] *= a
                else:
                    vx = (-b * c * particles_vz[index] * particles_vx[index] + b * d * particles_vy[index]) / v_ + a * particles_vx[index]
                    vy = (-b * c * particles_vz[index] * particles_vy[index] - b * d * particles_vx[index]) / v_ + a * particles_vy[index]
                    vz = b * c * v_ + a * particles_vz[index]
                    particles_vx[index], particles_vy[index], particles_vz[index] = vx, vy, vz
        return particlesEnergy, particles_vx, particles_vy, particles_vz, particles_detected
    
if __name__ == "__main__":
    fd = FinalDetector()
    fd.simulate()